"""
train_config.py
---------------
整合训练相关配置，形成统一的 TrainConfig 类。

功能：
- 数据集、模型、优化器、调度器、数据增强
- 损失函数（引用独立 loss_config.py 支持多任务 & 自定义损失）
- 评价指标（引用独立 metrics_config.py）
- 通用训练参数：epoch、日志、保存策略、分布式、混合精度等
- 支持从 YAML / JSON 文件加载配置
- 文件末尾包含测试代码
"""

from dataclasses import dataclass, field
from typing import Optional
import yaml
import json

# -----------------------------
# 子模块导入
# -----------------------------
from .dataset_config import DatasetConfig
from .augment_config import AugmentConfig
from .model_config import ModelConfig
from .optimizer_config import OptimizerManagerConfig
from .loss_config import LossConfig
from .metrics_config import MetricsConfig  # 引用独立文件
from .checkpoint_config import CheckpointConfig
from .logger_config import LoggerConfig
from .visualization_config import VisualizationConfig

# -----------------------------
# TrainConfig
# -----------------------------
@dataclass
class TrainConfig:
    dataset: DatasetConfig = None
    model: ModelConfig = None
    optimizer: OptimizerManagerConfig = None
    loss: LossConfig = field(default_factory=LossConfig)
    metrics: MetricsConfig = field(default_factory=MetricsConfig)  # 新增指标配置
    checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig)
    logger: LoggerConfig = field(default_factory=LoggerConfig)
    visualization: VisualizationConfig = field(default_factory=VisualizationConfig)

    # ----------------------
    # 通用训练参数
    # ----------------------
    num_epochs: int = 100
    batch_accumulate_steps: int = 1
    seed: int = 42
    device: str = "cuda"
    num_gpus: int = 1
    distributed: bool = False
    mixed_precision: bool = False
    gradient_clip: Optional[float] = 1.0

    # ----------------------
    # 评估与早停
    # ----------------------
    eval_interval: int = 1
    early_stopping_patience: Optional[int] = None

    # ----------------------
    # 方法：打印配置摘要
    # ----------------------
    def summary(self):
        print("=============== 🏋️‍♂️ 训练配置概览 ===============")
        print(f"📚 数据集: {self.dataset.dataset_name} ({getattr(self.dataset, 'dataset_type', 'builtin')})")
        print(f"🧠 模型: {self.model.model_name} ({getattr(self.model, 'model_type', 'builtin')})")
        print(f"⚙️ 优化器: {self.optimizer.optimizer.optimizer_type}")
        if self.optimizer.scheduler.scheduler_type:
            print(f"📉 调度器: {self.optimizer.scheduler.scheduler_type}")
        print(f"🧩 数据增强启用: {self.dataset.augment.image.use_augmentation}")
        self.loss.summary()
        self.metrics.summary()  # 打印指标信息
        print(f"训练设备: {self.device} (GPU数={self.num_gpus})")
        print(f"Epochs: {self.num_epochs}, Batch Accumulation: {self.batch_accumulate_steps}")
        print(f"混合精度: {self.mixed_precision}, 分布式: {self.distributed}")
        print(f"保存路径: {self.checkpoint.save_dir}, 日志路径: {self.logger.log_dir}")
        print("==================================================")

    # ----------------------
    # 从 YAML / JSON 文件加载
    # ----------------------
    @staticmethod
    def from_yaml(path: str) -> "TrainConfig":
        with open(path, "r", encoding="utf-8") as f:
            data = yaml.safe_load(f)
        return TrainConfig._from_dict(data)

    @staticmethod
    def from_json(path: str) -> "TrainConfig":
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        return TrainConfig._from_dict(data)

    @staticmethod
    def _from_dict(data: dict) -> "TrainConfig":
        """递归构造配置对象"""
        augment = AugmentConfig(**data.get("augment", {}))
        dataset = DatasetConfig(**data.get("dataset", {}),augment=augment)
        model = ModelConfig(**data.get("model", {}))
        optimizer = OptimizerManagerConfig(
            optimizer=data.get("optimizer", {}).get("optimizer"),
            scheduler=data.get("optimizer", {}).get("scheduler")
        )

        loss = LossConfig(**data.get("loss", {}))
        metrics = MetricsConfig(**data.get("metrics", {}))
        general = {k: v for k, v in data.items() if k not in ["dataset", "model", "optimizer", "augment", "loss", "metrics"]}
        return TrainConfig(
            dataset=dataset,
            model=model,
            optimizer=optimizer,
            loss=loss,
            metrics=metrics,
            **general
        )


# -----------------------------
# 测试代码
# -----------------------------
if __name__ == "__main__":
    print("=== 测试 TrainConfig ===")

    # 创建示例配置
    from .loss_config import SingleLossConfig

    train_config = TrainConfig(
        num_epochs=10,
        mixed_precision=True,
        loss=LossConfig(
            losses=[
                SingleLossConfig(loss_type="cross_entropy", weight=1.0),
                SingleLossConfig(loss_type="mse", weight=0.5)
            ]
        ),
        metrics=MetricsConfig(
            metrics=["accuracy", "f1"],
            custom_metrics={"custom_mse": lambda pred, target: ((pred - target)**2).mean()}
        )
    )

    # 打印配置摘要
    train_config.summary()
